#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Author:MC.Lee

# 依赖 pip install requests
# pip install 


import hashlib
import requests
import time
import base64
import string
import random
import struct
import socket
from Crypto.Cipher import AES
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET


class ierror:  
    WX_MsgCrypt_OK = {"code":0,"msg":"成功"}
    WX_MsgCrypt_ValidateSignature_Error = {"code":-40001,"msg":"验证签名错误"}
    WX_MsgCrypt_ParseXml_Error = {"code":-40002,"msg":"xml解析错误"}
    WX_MsgCrypt_ComputeSignature_Error = {"code":-40003,"msg":"计算签名错误"}
    WX_MsgCrypt_IllegalAesKey = {"code":-40004,"msg":"非法的AesKey"}
    WX_MsgCrypt_ValidateAppid_Error = {"code":-40005,"msg":"验证appid错误"}
    WX_MsgCrypt_EncryptAES_Error = {"code":-40006,"msg":"AES加密错误"}
    WX_MsgCrypt_DecryptAES_Error = {"code":-40007,"msg":"AES解密错误"}
    WX_MsgCrypt_IllegalBuffer = {"code":-40008,"msg":"非法的buffer"}
    WX_MsgCrypt_EncodeBase64_Error = {"code":-40009,"msg":"Base64加密错误"}
    WX_MsgCrypt_DecodeBase64_Error = {"code":-40010,"msg":"Base64解密错误"}
    WX_MsgCrypt_GenReturnXml_Error = {"code":-40011,"msg":"生成XML错误"}
    WX_MsgAdmin_XmlMsgType_Error = {"code":-40012,"msg":"XML消息类型错误"}
    WX_MsgAdmin_XmlMsgData_Error = {"code":-40013,"msg":"用于构建XML消息的数据错误"}


class WXPublic(object):

    def __init__(self, token, AESkey, MsgEncryptionType):
        '''
        @param token 
        @param AESkey 消息加密秘钥
        @param MsgEncryptionType 消息加密模式 Plaintext明文 compatible兼容模式 ciphertext密文
        '''
        self.__token = token
        self.__AESkey = base64.b64decode(AESKey + "=")
        self.__MsgEncryptionType = MsgEncryptionType
        # self.__MsgEncryptionType = "ciphertext"
        self.__AccessToken = {
            "access_token":None,
            "expires_at":0
        }
        self.appid = ""

    def check_signature(self, signature, timestamp, nonce, echostr):
        '''
        公众号接入验证接口
        '''
        check_list = [self.__token, timestamp, nonce]
        check_list.sort()
        sha1_str = hashlib.sha1()
        map(sha1_str.update, check_list)
        hashcode = sha1_str.hexdigest()
        print(f'验证签名：{hashcode}')
        if hashcode == signature:
            return ierror.WX_MsgCrypt_OK, echostr
        else:
            print(ierror.WX_MsgCrypt_ValidateSignature_Error['msg'])
            return ierror.WX_MsgCrypt_ValidateSignature_Error, ""

    def ComputemsgSignature(self, timestamp, nonce, encrypt):
        """计算公众平台的消息签名接口"""
        """用SHA1算法生成安全签名
        @param token:  票据
        @param timestamp: 时间戳
        @param encrypt: 密文
        @param nonce: 随机字符串
        @return: 安全签名
        """
        try:
            sortlist = [self.__token, timestamp, nonce, encrypt]
            sortlist.sort()
            sha = hashlib.sha1()
            sha.update("".join(sortlist))
            return  ierror.WX_MsgCrypt_OK, sha.hexdigest()
        except Exception as e:
            print(e, ierror.WX_MsgCrypt_ComputeSignature_Error['msg'])
            return  ierror.WX_MsgCrypt_ComputeSignature_Error, None

    def getAccessToken(self, appID, appSecret):
        '''
        从微信服务器获取accesstoken
        '''
        if checkAccessTokenExpores():
            return self.__AccessToken["access_token"]
        else:
            url = f"https://api.weixin.qq.com/cgi-bin/token?grant_type=client_credential&appid={appID}&secret={appSecret}"
            reponse = requests.get(url)
            data = reponse.json()
            if "access_token" in data:
                self.__AccessToken["access_token"] = data["access_token"]
                self.__AccessToken["expires_at"] = int(time.time()) + data["expires_in"]
                return self.__AccessToken["access_token"]
            else:
                print(f'错误代码：{data["errcode"]}，错误信息：{data["errmsg"]}')
                return None

    def checkAccessTokenExpores(self):
        '''
        检测access_token是否有效或过期
        过期时间剩余小于5分钟，默认过期
        '''
        if self.__AccessToken["access_token"] == None:
            return False
        elif self.__AccessToken["expires_at"] - 300 >= int(time.time()):
            return False
        else:
            return True

    def createMenu(self,data):
        '''
        创建菜单
        @param data 菜单json数据
        {"button":[...]} 一级菜单数组，最多3个菜单
        {"name":"菜单名称","sub_button":[...]} 二级菜单数组，个数应为1~5个
        {
            "type":"菜单按钮类型",
            "name":"菜单名称",
            "key":"click等点击类型必须,菜单KEY值，用于消息接口推送",
            "url":"view、miniprogram类型必须,网页 链接，用户点击菜单可打开链接，不超过1024字节。 type为miniprogram时，不支持小程序的老版本客户端将打开本url",
            "media_id":"media_id类型和view_limited类型必须,调用新增永久素材接口返回的合法media_id",
            "appid":"miniprogram类型必须", #小程序的appid（仅认证公众号可配置）
            "pagepath":"miniprogram类型必须", # 小程序的页面路径
        }
        @return {"errcode":0,"errmsg":"ok"}
        @tips: 一级菜单最多4个汉字，二级菜单最多7个汉字，多出来的部分将会以“...”代替
        @param type 1、click：点击推事件用户点击click类型按钮后，微信服务器会通过消息接口推送消息类型为event的结构给开发者（参考消息接口指南），并且带上按钮中开发者填写的key值，开发者可以通过自定义的key值与用户进行交互；
                    2、view：跳转URL用户点击view类型按钮后，微信客户端将会打开开发者在按钮中填写的网页URL，可与网页授权获取用户基本信息接口结合，获得用户基本信息。
                    3、scancode_push：扫码推事件用户点击按钮后，微信客户端将调起扫一扫工具，完成扫码操作后显示扫描结果（如果是URL，将进入URL），且会将扫码的结果传给开发者，开发者可以下发消息。
                    4、scancode_waitmsg：扫码推事件且弹出“消息接收中”提示框用户点击按钮后，微信客户端将调起扫一扫工具，完成扫码操作后，将扫码的结果传给开发者，同时收起扫一扫工具，然后弹出“消息接收中”提示框，随后可能会收到开发者下发的消息。
                    5、pic_sysphoto：弹出系统拍照发图用户点击按钮后，微信客户端将调起系统相机，完成拍照操作后，会将拍摄的相片发送给开发者，并推送事件给开发者，同时收起系统相机，随后可能会收到开发者下发的消息。
                    6、pic_photo_or_album：弹出拍照或者相册发图用户点击按钮后，微信客户端将弹出选择器供用户选择“拍照”或者“从手机相册选择”。用户选择后即走其他两种流程。
                    7、pic_weixin：弹出微信相册发图器用户点击按钮后，微信客户端将调起微信相册，完成选择操作后，将选择的相片发送给开发者的服务器，并推送事件给开发者，同时收起相册，随后可能会收到开发者下发的消息。
                    8、location_select：弹出地理位置选择器用户点击按钮后，微信客户端将调起地理位置选择工具，完成选择操作后，将选择的地理位置发送给开发者的服务器，同时收起位置选择工具，随后可能会收到开发者下发的消息。
                    9、media_id：下发消息（除文本消息）用户点击media_id类型按钮后，微信服务器会将开发者填写的永久素材id对应的素材下发给用户，永久素材类型可以是图片、音频、视频、图文消息。请注意：永久素材id必须是在“素材管理/新增永久素材”接口上传后获得的合法id。
                    10、view_limited：跳转图文消息URL用户点击view_limited类型按钮后，微信客户端将打开开发者在按钮中填写的永久素材id对应的图文消息URL，永久素材类型只支持图文消息。请注意：永久素材id必须是在“素材管理/新增永久素材”接口上传后获得的合法id
        '''
        url =  f"https://api.weixin.qq.com/cgi-bin/menu/create?access_token={self.getAccessToken()}"
        response = requests.post(url,data)
        res = reponse.json()
        print(f'错误代码：{res["errcode"]}，错误信息：{res["errmsg"]}')
        return res

    def getMenu(self):
        '''
        获取公众号当前菜单json数据
        @return 成功返回json数据，失败返回错误代码json
        '''
        url = f"https://api.weixin.qq.com/cgi-bin/menu/get?access_token={self.getAccessToken()}"
        response = requests.get(url)
        res = reponse.json()
        if "menu" in res:
            print('菜单获取成功')
        else:
            print(f'错误代码：{res["errcode"]}，错误信息：{res["errmsg"]}')
        return res

    def delMenu(self):
        '''
        删除公众号当前菜单
        @return 成功返回{"errcode":0,"errmsg":"ok"}，失败返回错误代码
        '''
        url = f"https://api.weixin.qq.com/cgi-bin/menu/delete?access_token={self.getAccessToken()}"
        response = requests.get(url)
        res = reponse.json()
        if res["errcode"] == 0:
            print('菜单删除成功')
        else:
            print(f'错误代码：{res["errcode"]}，错误信息：{res["errmsg"]}')
        return res

    def parseXmlMsg(self, xmlData):
        '''
        解析xml消息，返回消息字典
        @param xmlData 获取的xml消息
        @return 消息字典
        '''
        data = {}
        try:
            root = ET.fromstring(xmlStr)
            for child in root:
                data[child.tag] = child.text
            # print(data)
            return ierror.WX_MsgCrypt_OK, data
        except Exception as e:
            print(e, ierror.WX_MsgCrypt_ParseXml_Error['msg'])
            return ierror.WX_MsgCrypt_ParseXml_Error, None


    def textXmlMsg_Template(self):
        temp = '''<xml>
            <ToUserName><![CDATA[%(tousername)s]]></ToUserName>
            <FromUserName><![CDATA[%(fromusername)s]]></FromUserName>
            <CreateTime>%(createtime)s</CreateTime>
            <MsgType><![CDATA[text]]></MsgType>
            <Content><![CDATA[%(content)s]]></Content>
        </xml>
        '''
        return temp

    def imageXmlMsg_Template(self):
        temp = '''<xml>
        <ToUserName><![CDATA[%(tousername)s]]></ToUserName>
        <FromUserName><![CDATA[%(fromusername)s]]></FromUserName>
        <CreateTime>%(createtime)s</CreateTime>
        <MsgType><![CDATA[image]]></MsgType>
        <Image>
            <MediaId><![CDATA[%(mediaid)s]]></MediaId>
        </Image>
        </xml>
        '''
        return temp

    def voiceXmlMsg_Template(self):
        temp = '''<xml>
        <ToUserName><![CDATA[%(tousername)s]]></ToUserName>
        <FromUserName><![CDATA[%(fromusername)s]]></FromUserName>
        <CreateTime>%(createtime)s</CreateTime>
        <MsgType><![CDATA[voice]]></MsgType>
        <Voice>
          <MediaId><![CDATA[%(mediaid)s]]></MediaId>
        </Voice>
      </xml>'''
        return temp

    def videoXmlMsg_Template(self, msgData):
        temp = '''<xml>
        <ToUserName><![CDATA[%(tousername)s]]></ToUserName>
        <FromUserName><![CDATA[%(fromusername)s]]></FromUserName>
        <CreateTime>%(createtime)s</CreateTime>
        <MsgType><![CDATA[video]]></MsgType>
        <Video>
          <MediaId><![CDATA[%(mediaid)s]]></MediaId>
          '''
        if "title" in msgData:
            temp = temp + '''<Title><![CDATA[%(title)s]]></Title>'''
        if "description" in msgData:
            temp = temp + '''<Description><![CDATA[%(description)s]]></Description>'''
        temp = temp + '''</Video>
        </xml>'''
        return temp

    def musicXmlMsg_Template(self, msgData):
        temp = '''<xml>
        <ToUserName><![CDATA[%(tousername)s]]></ToUserName>
        <FromUserName><![CDATA[%(fromusername)s]]></FromUserName>
        <CreateTime>%(createtime)s</CreateTime>
        <MsgType><![CDATA[music]]></MsgType>
        <Music>
          '''
        if "title" in msgData:
            # 音乐标题
            temp = temp + '''<Title><![CDATA[%(title)s]]></Title>'''
        if "description" in msgData:
            # 音乐描述
            temp = temp + '''<Description><![CDATA[%(description)s]]></Description>'''
        if "musicurl" in msgData:
            # 音乐链接
            temp = temp + '''<MusicUrl><![CDATA[%(musicurl)s]]></MusicUrl>'''
        if "HQmusicurl" in msgData:
            # 高质量音乐链接，WIFI环境优先使用该链接播放音乐
            temp = temp + '''<HQMusicUrl><![CDATA[%(HQmusicurl)s]]></HQMusicUrl>'''
        # ThumbMediaId	缩略图的媒体id，通过素材管理中的接口上传多媒体文件，得到的id
        temp = temp + '''<ThumbMediaId><![CDATA[%(mediaid)s]]></ThumbMediaId>
        </Music>
      </xml>'''
        return temp

    def articlesXmlMsg_template(self,msgData):
        try:
            temp = f'''<xml>
            <ToUserName><![CDATA[{msgData["tousername"]}]]></ToUserName>
            <FromUserName><![CDATA[{msgData["fromusername"]}]]></FromUserName>
            <CreateTime>{msgData["createtime"]}</CreateTime>
            <MsgType><![CDATA[news]]></MsgType>
            <ArticleCount>{len(msgData["articles"])}</ArticleCount>
            <Articles>
            '''
            # articles 为列表类型，存储每篇图文
            for item in msgData["articles"]:
                temp += f'''<item>
                <Title><![CDATA[{item["title"]}]]></Title>
                <Description><![CDATA[{item["description"]}]]></Description>
                <PicUrl><![CDATA[{item["picurl"]}]]></PicUrl>
                <Url><![CDATA[{item["url"]}]]></Url>
            </item>'''
            temp += '''</Articles>
            </xml>'''
            return ierror.WX_MsgCrypt_OK, temp
        except Exception as e:
            print(e, ierror.WX_MsgAdmin_XmlMsgData_Error['msg'])
            return  ierror.WX_MsgAdmin_XmlMsgData_Error, None

    def createXmlMsg(self, msgType, msgData):
        '''
        根据消息类型，创建对应的XML格式消息
        @param msgType 消息类型：text image voice video music articles
        @param msgData 用于构建XML消息的数据，字典格式
        '''
        if msgType == 'text':
            xml_template = self.textXmlMsg_Template()
        elif msgType == 'image':
            xml_template = self.imageXmlMsg_Template()
        elif msgType == "voice":
            xml_template = self.voiceXmlMsg_Template()
        elif msgType == "video":
            # xml消息中包含可选参数，需要传入data构建模板
            xml_template = self.videoXmlMsg_Template(msgData)
        elif msgType == "music":
            # xml消息中包含可选参数，需要传入data构建模板
            xml_template = self.musicXmlMsg_Template(msgData)
        elif msgType == "articles":
            # 图文消息可能包含多个图文，直接返回生成的XML
            return self.articlesXmlMsg_template(msgData)
        else:
            print(ierror.WX_MsgAdmin_XmlMsgType_Error['msg'])
            return  ierror.WX_MsgAdmin_XmlMsgType_Error, None
        try:
            res = xml_template % msgData
            return ierror.WX_MsgCrypt_OK, res
        except Exception as e:
            print(e, ierror.WX_MsgAdmin_XmlMsgData_Error['msg'])
            return ierror.WX_MsgAdmin_XmlMsgData_Error, None

    def replyMsg(self, xmlMsg, nonce=None, timestamp=None):
        '''
        根据加密类型回复消息
        Plaintext明文 compatible兼容模式 ciphertext密文
        '''
        if self.__MsgEncryptionType == "plaintext":
            return xmlMsg
        elif self.__MsgEncryptionType == "ciphertext":
            return self.encryptMsg(xmlMsg, nonce, timestamp)
        else:
            pass

    def encryptMsg(self, xmlMsg, nonce, timestamp=None):
        '''
        #将公众号回复用户的消息加密打包
        @param xmlMsg: 企业号待回复用户的消息，xml格式的字符串
        @param TimeStamp: 时间戳，可以自己生成，也可以用URL参数的timestamp,如为None则自动用当前时间
        @param Nonce: 随机串，可以自己生成，也可以用URL参数的nonce
        sEncryptMsg: 加密后的可以直接回复用户的密文，包括msg_signature, timestamp, nonce, encrypt的xml格式的字符串,
        @return 成功0，sEncryptMsg,失败返回对应的错误码None
        '''
        pc = Prpcrypt(self.__AESkey)
        ret, encrypt = pc.encrypt(xmlMsg,self.appid)
        if ret != 0:
            return ret, None
        if timestamp is None:
            timestamp = str(int(time.time()))
        # 生成安全签名
        ret, signature = self.ComputemsgSignature(timestamp, nonce, encrypt)
        if ret != 0:
            return ret, None
        return XMLParse.generate(encrypt, signature, timestamp, nonce)


class XMLParse:
    """提供提取消息格式中的密文及生成回复消息格式的接口"""

    # xml消息模板
    AES_TEXT_RESPONSE_TEMPLATE = """<xml>
<Encrypt><![CDATA[%(msg_encrypt)s]]></Encrypt>
<MsgSignature><![CDATA[%(msg_signaturet)s]]></MsgSignature>
<TimeStamp>%(timestamp)s</TimeStamp>
<Nonce><![CDATA[%(nonce)s]]></Nonce>
</xml>"""

    def extract(self, xmltext):
        """提取出xml数据包中的加密消息
        @param xmltext: 待提取的xml字符串
        @return: 提取出的加密消息字符串
        """
        try:
            xml_tree = ET.fromstring(xmltext)
            encrypt  = xml_tree.find("Encrypt")
            touser_name    = xml_tree.find("ToUserName")
            return  ierror.WX_MsgCrypt_OK, encrypt.text, touser_name.text
        except Exception as e:
            print(e, ierror.WX_MsgCrypt_ParseXml_Error['msg'])
            return  ierror.WX_MsgCrypt_ParseXml_Error,None,None
    
    def generate(self, encrypt, signature, timestamp, nonce):
        """生成xml消息
        @param encrypt: 加密后的消息密文
        @param signature: 安全签名
        @param timestamp: 时间戳
        @param nonce: 随机字符串
        @return: 生成的xml字符串
        """
        resp_dict = {
                    'msg_encrypt' : encrypt,
                    'msg_signaturet': signature,
                    'timestamp'    : timestamp,
                    'nonce'        : nonce,
                        }
        resp_xml = self.AES_TEXT_RESPONSE_TEMPLATE % resp_dict
        return resp_xml

class PKCS7Encoder():
    """提供基于PKCS7算法的加解密接口"""

    block_size = 32
    def encode(self, text):
        """ 对需要加密的明文进行填充补位
        @param text: 需要进行填充补位操作的明文
        @return: 补齐明文字符串
        """
        text_length = len(text)
        # 计算需要填充的位数
        amount_to_pad = self.block_size - (text_length % self.block_size)
        if amount_to_pad == 0:
            amount_to_pad = self.block_size
        # 获得补位所用的字符
        pad = chr(amount_to_pad).encode()
        return text + pad * amount_to_pad

    def decode(self, decrypted):
        """删除解密后明文的补位字符
        @param decrypted: 解密后的明文
        @return: 删除补位字符后的明文
        """
        pad = ord(decrypted[-1])
        if pad<1 or pad >32:
            pad = 0
        return decrypted[:-pad]    


class Prpcrypt(object):
    """提供接收和推送给公众平台消息的加解密接口"""

    def __init__(self,key):
        #self.key = base64.b64decode(key+"=")
        self.key = key
        # 设置加解密模式为AES的CBC模式
        self.mode = AES.MODE_CBC


    def encrypt(self,text,appid):
        """对明文进行加密
        @param text: 需要加密的明文
        @return: 加密得到的字符串
        """
        # 16位随机字符串添加到明文开头
        len_str = struct.pack("I", socket.htonl(len(text.encode())))
        # text = self.get_random_str() + binascii.b2a_hex(len_str).decode() + text + appid
        text = self.get_random_str().encode() + len_str + text.encode() + appid.encode()
        # 使用自定义的填充方式对明文进行补位填充
        pkcs7 = PKCS7Encoder()
        text = pkcs7.encode(text)
        # 加密
        cryptor = AES.new(self.key,self.mode,self.key[:16])
        try:
            ciphertext = cryptor.encrypt(text)
            # 使用BASE64对加密后的字符串进行编码
            return ierror.WX_MsgCrypt_OK, base64.b64encode(ciphertext)
        except Exception as e:
            print(e,ierror.WX_MsgCrypt_EncryptAES_Error["msg"])
            return  ierror.WX_MsgCrypt_EncryptAES_Error,None

    def decrypt(self,text,appid):
        """对解密后的明文进行补位删除
        @param text: 密文
        @return: 删除填充补位后的明文
        """
        try:
            cryptor = AES.new(self.key,self.mode,self.key[:16])
            # 使用BASE64对密文进行解码，然后AES-CBC解密
            plain_text  = cryptor.decrypt(base64.b64decode(text))
        except Exception as e:
            print(e, ierror.WX_MsgCrypt_DecryptAES_Error['msg'])
            return  ierror.WX_MsgCrypt_DecryptAES_Error,None
        try:
            pad = ord(plain_text[-1])
            # 去掉补位字符串
            #pkcs7 = PKCS7Encoder()
            #plain_text = pkcs7.encode(plain_text)
            # 去除16位随机字符串
            content = plain_text[16:-pad]
            xml_len = socket.ntohl(struct.unpack("I",content[ : 4])[0])
            xml_content = content[4 : xml_len+4]
            from_appid = content[xml_len+4:]
        except Exception as e:
            print(e, ierror.WX_MsgCrypt_IllegalBuffe['msg'])
            return  ierror.WX_MsgCrypt_IllegalBuffer,None
        if  from_appid != appid:
            print(ierror.WX_MsgCrypt_ValidateAppid['msg'])
            return ierror.WX_MsgCrypt_ValidateAppid_Error,None
        return 0,xml_content

    def get_random_str(self):
        """ 随机生成16位字符串
        @return: 16位字符串
        """
        rule = string.ascii_letters + string.digits
        str = random.sample(rule, 16)
        return "".join(str)

if __name__ == "__main__":
    data = {
        'tousername':"123",
        "fromusername":"456",
        "createtime":int(time.time()),
        "articles":[
            {
                "title":"图文一",
                "description":"图文测试",
                "picurl":"http:www.baidu.com",
                "url":"www.baidu.com"
            },
            {
                "title":"图文er",
                "description":"图文测试",
                "picurl":"http:www.baidu.com",
                "url":"www.baidu.com"
            },
        ],
    }
    wx = WXPublic('123','abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG',"ciphertext")
    ret = wx.createXmlMsg("articles", data)
    nonce = Prpcrypt("key").get_random_str()
    ret = wx.replyMsg(ret[1],nonce)
    print(ret)
    